import json

def gen_prompt_ner_ds(data_path, save_path):
    list_data_dict = []
    with open(data_path, "r", encoding='utf-8') as f:
        list_data_dict.extend(json.load(f))

    ds = []

    for exp in list_data_dict:
        prompt = f"请抽取出给定句子中的所有命名实体二元组。给定句子：{exp['sentence'].lower()}"
        response = f"{[(ele['e_span'], ele['e_type']) for ele in exp['entities']]}".lower().replace("'", "")
        ds.append({"prompt": prompt, "response": response})
    with open(save_path, "w", encoding='utf-8') as f:
        json.dump(ds, f, ensure_ascii=False, indent=4)

    return ds

def gen_prompt_ee_ds(data_path, save_path):
    list_data_dict = []
    with open(data_path, "r", encoding='utf-8') as f:
        list_data_dict.extend(json.load(f))

    ds = []

    for exp in list_data_dict:
        prompt = f"请抽取出给定句子中的所有事件三元组。给定句子：{exp['sentence'].lower()}"
        response = f"{[(ele['e_span'], ele['e_type'], ele['e_role']) for ele in exp['events']]}".lower().replace("'", "")
        ds.append({"prompt": prompt, "response": response})
    with open(save_path, "w", encoding='utf-8') as f:
        json.dump(ds, f, ensure_ascii=False, indent=4)

    return ds

#%%
with open("./data/ace_ner_train.json", "r", encoding='utf-8') as f:
    js = json.load(f)

print(js[0]["prompt"], "\n", js[0]["response"])

print(js[0]["response"][2])
#%%
import os
paths = os.listdir("./ds")
ds_path = []
op_path = []
for path in paths:
    ds = path.split("_")[0].lower()
    task = path.split("_")[-1].lower()

    ds_path.append(f"./ds/{path}/train.json")

    op_path.append(f"./data/{ds}_{task}_train.json")

for i,o in zip(ds_path, op_path):
    gen_prompt_ner_ds(i, o)
#%%
import os
paths = os.listdir("./ds/")
paths
#%%
